import requests

# 基础初始化设置
base_url = "http://172.26.167.51:11434/api"
headers = {
    "Content-Type": "application/json"
}
model_name = "qwen2.5"


def generate_completion(prompt, model=model_name):
    url = f"{base_url}/generate"
    data = {
        "model": model,
        "prompt": prompt,
        "stream": False
    }
    response = requests.post(url, headers=headers, json=data)
    return response.json().get('response', '')


def generate_completion_stream(prompt, model=model_name):
    url = f"{base_url}/generate"
    data = {
        "model": model,
        "prompt": prompt,
        "stream": True
    }
    response = requests.post(url, headers=headers, json=data, stream=True)
    result = ""
    for line in response.iter_lines():
        if line:
            result += line.decode('utf-8')
    return result


# 示例调用
stream_completion = generate_completion_stream("讲解机器学习的应用。")
print("流式生成文本补全:", stream_completion)

# # 示例调用
# completion = generate_completion("介绍一下人工智能。")
# print("生成文本补全:", completion)
